import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl

x1 = np.array([[0, 0], [2, 1], [1, 0]])
x2 = np.array([[-1, 1], [-2, 0], [-2, -1]])
x3 = np.array([[0, -2], [0, -1], [1, -2]])
y1 = np.array([1]*len(x1))
y2 = np.array([2]*len(x2))
y3 = np.array([3]*len(x3))

X = np.concatenate((x1,x2,x3),axis=0)
y = np.concatenate((y1,y2,y3), axis=0)
print("训练样本的特征：\n",X)
print("训练样本的标签:", y)

eps = 1e-6
labels = np.unique(y)
K = len(labels)
N, M = X.shape

parameters = {
    "mean": np.zeros((K, M)),  # shape: (K, M)
    "sigma": np.zeros((K, M)),  # shape: (K, M)
    "prior": np.zeros((K,)),  # shape: (K,)
}

for i, c in enumerate(labels):
    X_c = X[y == c, :]

    parameters["mean"][i, :] = np.mean(X_c, axis=0)
    parameters["sigma"][i, :] = np.var(X_c, axis=0) + eps
    parameters["prior"][i] = X_c.shape[0] / N

def predict(x, labels):
    K = len(labels)
    log_posterior = np.zeros((x.shape[0], K))
    for i in range(K):
        mu = parameters["mean"][i]
        prior = parameters["prior"][i]
        sigsq = parameters["sigma"][i]

        # log likelihood = log X | N(mu, sigsq)
        log_likelihood = -0.5 * np.sum(np.log(2 * np.pi * sigsq))
        log_likelihood -= 0.5 * np.sum(((x - mu) ** 2) / sigsq, axis=1)

        log_posterior[:, i] = log_likelihood + np.log(prior)
    pred = labels[log_posterior.argmax(axis=1)]
    return pred, log_posterior

p = np.array([[-2, 2],])
pred, log_posterior= predict(p, labels)
print(f"样本属于第{pred}类", "，对数概率为：", log_posterior)

# 画图
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1  # 第0列的范围
y_min, y_max = X[:, 1].min() -1, X[:, 1].max() + 1  # 第1列的范围
xx, yy = np.meshgrid(np.linspace(x_min, x_max, 400), np.linspace(y_min, y_max, 400))  # 生成网格采样点
grid_test = np.stack((xx.flat, yy.flat), axis=1)  # 测试点  （xx.flat降维）
y_predict, _ = predict(grid_test, labels)

cm_pt = mpl.colors.ListedColormap(['r', 'g', 'b'])  # 样本点颜色（样本分为3个类，三个颜色）
cm_bg = mpl.colors.ListedColormap(['b', 'y', 'gray'])  # 背景颜色
mpl.rcParams['font.sans-serif'] = [u'SimHei']   # 设置字体为SimHei显示中文
mpl.rcParams['axes.unicode_minus'] = False  # 设置正常显示字符
plt.xlim(x_min, x_max)
plt.ylim(y_min, y_max)  # 设置坐标范围
plt.pcolormesh(xx, yy, y_predict.reshape(xx.shape), shading='auto', cmap=cm_bg)  # 绘制网格背景
plt.scatter(X[:, 0], X[:, 1], c=y, cmap=cm_pt, marker='o')  # 绘制样本点
plt.title(u'Bayes分类', fontsize=15)
plt.show()
